import torch
import torch.nn as nn
import torch.nn.functional as F

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Neural Network using device: {device}")

class PokerNet(nn.Module):
    """德州扑克AI的神经网络模型"""
    
    def __init__(self, input_dim: int, hidden_dim: int = 256, num_actions: int = 4):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim).to(device)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim).to(device)
        self.fc3 = nn.Linear(hidden_dim, hidden_dim).to(device)
        self.action_head = nn.Linear(hidden_dim, num_actions).to(device)  # 动作类型
        self.value_head = nn.Linear(hidden_dim, 1).to(device)  # 状态价值
        self.bet_head = nn.Linear(hidden_dim, 1).to(device)  # 下注金额
        
    def forward(self, x):
        x = x.to(device)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        
        action_probs = F.softmax(self.action_head(x), dim=-1)
        state_value = self.value_head(x)
        bet_amount = F.relu(self.bet_head(x))  # 确保下注金额非负
        
        return action_probs, state_value, bet_amount 